package com.xmzs.common.gpt.handler;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson2.JSONObject;
import com.xmzs.common.core.domain.model.LoginUser;
import com.xmzs.common.core.exception.base.BaseException;
import com.xmzs.common.core.utils.SpringUtils;
import com.xmzs.common.core.utils.StringUtils;
import com.xmzs.common.gpt.config.LocalCache;
import com.xmzs.common.gpt.entity.chat.ChatCompletion;
import com.xmzs.common.gpt.holder.WebSocketSessionHolder;
import com.xmzs.common.gpt.listener.WebSocketEventListener;
import com.xmzs.common.gpt.openai.OpenAiStreamClient;
import com.xmzs.common.gpt.entity.chat.Message;
import com.xmzs.common.gpt.service.TextReviewService;
import com.xmzs.common.gpt.utils.WebSocketUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import static com.xmzs.common.gpt.constant.WebSocketConstants.LOGIN_USER_KEY;

/**
 * WebSocketHandler 实现类
 *
 * @author zendwang
 */
@Slf4j
public class PlusWebSocketHandler extends AbstractWebSocketHandler {

    /**
     * 上下文长度
     */
    private final int contextSize = 10;

    /**
     * 审核状态 1 代表合法
     */
    private static final String conclusionType = "1";

    private OpenAiStreamClient openAiStreamClient;

    private TextReviewService textReviewService;

    /**
     * 连接成功后
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
        WebSocketSessionHolder.addSession(loginUser.getUserId(), session);
        log.info("[connect] sessionId: {},userId:{},userType:{}", session.getId(), loginUser.getUserId(), loginUser.getUserType());
    }

    /**
     * 处理发送来的文本消息
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
        log.info("PlusWebSocketHandler, 连接：" + session.getId() + "，已收到消息:" + message.getPayload());
        // 判断文本是否合规
//        textReviewService = (TextReviewService) SpringUtils.context().getBean("textReviewService");
//        String type = textReviewService.textReview(message.getPayload());
//        if (!conclusionType.equals(type) && StringUtils.isNotEmpty(type)) {
//            HashMap<Object, Object> msgMap = new HashMap<>(10);
//            msgMap.put("content", "文本不合规,请修改!");
//            String jsonStr = JSONUtil.toJsonStr(msgMap);
//            WebSocketUtils.sendMessage(session, jsonStr);
//            WebSocketUtils.sendMessage(session, "[DONE]");
//            return;
//        }
        JSONObject jsonObject = JSONObject.parseObject(message.getPayload());
        String systemMsg = getValue(jsonObject,"systemMessage","");
        // 是否需要携带上下文信息
        Boolean isContext = Boolean.valueOf(getValue(jsonObject, "isContext", "true"));
        WebSocketEventListener eventSourceListener = new WebSocketEventListener(session);
        String messageContext = (String) LocalCache.CACHE.get(String.valueOf(loginUser.getUserId()));
        List<Message> messages = new ArrayList<>();
        if (StrUtil.isNotBlank(messageContext)) {
            messages = JSONUtil.toList(messageContext, Message.class);
            if (messages.size() >= contextSize) {
                messages = messages.subList(1, contextSize);
            }
            // 清空上下文信息
            if(!isContext){
                messages.clear();
            }
            Message currentMessage = Message.builder().content(message.getPayload()).role(Message.Role.USER).build();
            messages.add(currentMessage);
        } else {
            // 判断是否设置了系统角色
            if(StrUtil.isNotEmpty(systemMsg)){
                Message systemMessage = Message.builder().content(systemMsg).role(Message.Role.SYSTEM).build();
                messages.add(systemMessage);
            }
            Message currentMessage = Message.builder().content(getValue(jsonObject,"prompt","你好")).role(Message.Role.USER).build();
            messages.add(currentMessage);
        }
        ChatCompletion chatCompletion = ChatCompletion
            .builder()
            .model(getValue(jsonObject,"model",ChatCompletion.Model.GPT_3_5_TURBO_16K.getName()))
            .messages(messages)
            .temperature(Double.parseDouble(getValue(jsonObject,"temperature","0.2")))
            .stream(Boolean.valueOf(getValue(jsonObject,"stream","true")))
            .build();
        openAiStreamClient = (OpenAiStreamClient) SpringUtils.context().getBean("openAiStreamClient");
        openAiStreamClient.streamChatCompletion(chatCompletion, eventSourceListener);
        LocalCache.CACHE.put(String.valueOf(loginUser.getUserId()), JSONUtil.toJsonStr(messages), LocalCache.TIMEOUT);
    }

    /**
     * 根据key获取Value值
     *
     * @Date 2023/7/27
     * @param jsonObject
     * @param key
     * @param defaultValue
     * @return String
     **/
    public String getValue(JSONObject jsonObject,String key,String defaultValue){
        String value = (String)jsonObject.get(key);
        if(StrUtil.isEmpty(value)){
            return defaultValue;
        }
        return value;
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        super.handleBinaryMessage(session, message);
    }

    /**
     * 心跳监测的回复
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        WebSocketUtils.sendPongMessage(session);
    }

    /**
     * 连接出错时
     *
     * @param session
     * @param exception
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("[transport error] sessionId: {} , exception:{}", session.getId(), exception.getMessage());
    }

    /**
     * 连接关闭后
     *
     * @param session
     * @param status
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        LoginUser loginUser = (LoginUser) session.getAttributes().get(LOGIN_USER_KEY);
        WebSocketSessionHolder.removeSession(loginUser.getUserId());
        log.info("[disconnect] sessionId: {},userId:{},userType:{}", session.getId(), loginUser.getUserId(), loginUser.getUserType());
    }

    /**
     * 是否支持分片消息
     *
     * @return
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
